from scipy.optimize import minimize
import numpy as np


def cal_weight_nwp_flow(index, optical_flow, pwafs_mat, true_mat):
    X = np.column_stack([pwafs_mat[:,index].ravel(), optical_flow[:,index].ravel()])
    Y = true_mat[:,index].ravel()
    # 定义目标函数（MSE），这里只有一个变量 w1
    def objective(w1):
        w2 = 1 - w1
        prediction = w1 * X[:, 0] + w2 * X[:, 1]
        mse = ((prediction - Y) ** 2).mean()
        return mse
    # 初始猜测
    initial_guess = [0.5]

    # 优化
    result = minimize(objective, initial_guess)

    # 最优权重
    w1_optimal = result.x[0]
    w2_optimal = 1 - w1_optimal

    print("Optimal w1:", w1_optimal)
    print("Optimal w2:", w2_optimal)

    return w1_optimal, w2_optimal